import os, shutil
import skimage.io as io
import os.path as osp
import json
from tqdm import tqdm
from tool import filesystem, via_tool # export PYTHONPATH=$PYTHONPATH:`pwd`
from pycocotools import coco
from collections import defaultdict
import matplotlib
import matplotlib.pyplot as plt


def coco_to_via(aau_rain_snow_dir, 
                save_dir, 
                sub_dir="aau_rain_snow", 
                ann_file="aauRainSnow-rgb.json",
                via_name="via_region_data.json"):
    """
    aau rain snow tool
    将此数据集转换成via格式
    """
    rainSnowRgbGt = coco.COCO(osp.join(aau_rain_snow_dir, ann_file))
    label_dict = rainSnowRgbGt.cats
    img_group = defaultdict(list)
    for k, img in rainSnowRgbGt.imgs.items():
        sub_path = img["file_name"]
        # 按名字分组
        img_group[osp.dirname(sub_path)].append(img)

    save_path_dir = osp.join(save_dir, sub_dir)
    os.makedirs(save_path_dir, exist_ok=True)
    total_json = dict()
    # count_dict = defaultdict(int)
    for k,imgs in img_group.items():

        for img in tqdm(imgs):
            save_name = img["file_name"].replace("/", "_")
            file_path = osp.join(aau_rain_snow_dir, img["file_name"])
            file_size = osp.getsize(file_path)

            one_json = dict()
            one_json["filename"] = save_name
            one_json["size"] = file_size
            one_json["file_attributes"] = {}

            annIds = rainSnowRgbGt.getAnnIds(imgIds=[img["id"]])
            anns = rainSnowRgbGt.loadAnns(annIds)
            regions = []
            for ann in anns:
                region = dict()
                # if len(ann["segmentation"]) > 1: print(file_name)
                # count_dict[len(ann["segmentation"])] += 1
                label = label_dict[ann["category_id"]]["name"]
                shape_attributes = {
                    "name": "rect",
                    "x": ann["bbox"][0],
                    "y": ann["bbox"][1],
                    "width": ann["bbox"][2],
                    "height": ann["bbox"][3]
                }
                region_attributes = {
                    "original_grayscale_value": ann["original_grayscale_value"],
                    "iscrowd": ann["iscrowd"],
                    "label": label
                }
                region["shape_attributes"] = shape_attributes
                region["region_attributes"] = region_attributes
                regions.append(region)
            one_json["regions"] = regions
            total_json[save_name + str(file_size)] = one_json
            shutil.copy(file_path, osp.join(save_path_dir, save_name))

    # convert
    # print(count_dict)
    with open(osp.join(save_path_dir, via_name), "w") as wf:
        wf.write(json.dumps(total_json))


def aau_show(aau_rain_snow_dir, save_dir, ann_file="aauRainSnow-rgb.json"):
    rainSnowRgbGt = coco.COCO(osp.join(aau_rain_snow_dir, ann_file))

    for i in range(1950, 1951):
        chosenImgId = i
        annIds = rainSnowRgbGt.getAnnIds(imgIds=[chosenImgId])
        anns = rainSnowRgbGt.loadAnns(annIds)

        rgbImg = rainSnowRgbGt.loadImgs([chosenImgId])[0]

        print('Found ' + str(len(anns)) + ' annotations at image ID ' + str(chosenImgId) + '. Image file: ' + rgbImg['file_name'])

        for ann in anns:
            print('Annotation #' + str(ann['id']) + ': ' + rainSnowRgbGt.loadCats(ann['category_id'])[0]['name'])

        matplotlib.rcParams['interactive'] == False
        print("\nRGB Image")
        I = io.imread(osp.join(aau_rain_snow_dir, rgbImg['file_name']))
        plt.gcf().clear()
        plt.axis('off')
        plt.imshow(I)
        rainSnowRgbGt.showAnns(anns)
        plt.show()

        # For some reason, the image won't show in some Windows/Anaconda configurations. If this is the case, print the image instead
        plt.savefig("./rgb-" + str(chosenImgId).zfill(5) + ".png")

        # print("\nThermal Image")
        # # Load thermal annotations
        # I = io.imread('./' + thermalImg['file_name'])
        # plt.gcf().clear()
        # plt.axis('off')
        # plt.imshow(I);
        # rainSnowThermalGt.showAnns(thermalAnns)
        # plt.show()
        # plt.savefig("Samples/thermal-" + str(chosenImgId).zfill(5) + ".png")


if __name__ == "__main__":
    # # car truck bus
    # image_dir = "/home/xc/work/data/car/bdd100k/bdd100k/images/100k"
    # save_dir = "/home/xc/work/data/car/train/bdd100"
    # label_files = [
    #     "/home/xc/work/data/car/bdd100k/bdd100k/labels/det_20/det_train.json",
    #     "/home/xc/work/data/car/bdd100k/bdd100k/labels/det_20/det_val.json"
    # ]

    # json_to_via(image_dir, save_dir, label_files)


    # data_ori_path = r"/home/xc/work/data/car/train/bdd100/79/via_region_data_ori.json"
    # via_path =      r"/home/xc/work/data/car/train/bdd100/79/via_region_data_size.json"
    # via_tool.convert_to_via(data_ori_path, via_path)




    aau_rain_snow_dir = "/home/xc/work/data/car/AAU_RainSnow_Traffic_Surveillance_Dataset" 
    save_dir = "/home/xc/work/data/car/train"
    # 将所有图片合并到一个文件夹
    coco_to_via(aau_rain_snow_dir, save_dir)

    # aau_show(aau_rain_snow_dir, save_dir)

    
